import numpy as np
import pandas as pd
import tushare as ts
import config
import datetime

def getdata(det):
    pro = ts.pro_api()

    positive = pd.DataFrame(columns=['symbol','name','open','close','trade_date'])
    negative = pd.DataFrame(columns=['symbol','name','open','close','trade_date'])
    positive['symbol'].astype(str)
    negative['symbol'].astype(str)

    stocklist = pro.stock_basic(exchange='', list_status='L', market='主板',fields='ts_code,symbol,name')

    stocks = stocklist
    num = 0
    symbols = []
    trade_date = ''

    for stock in stocklist.itertuples():
        ts_code = getattr(stock,'ts_code')
        symbols.append(ts_code.split('.')[0])
        num = num + 1
        if num == 50:
            #获取实时行情数据
            datas = ts.get_realtime_quotes(symbols)

            print(datas.head())
            num = 0
            symbols.clear()

            for df in datas.itertuples():
                name = getattr(df, 'name')
                symbol = getattr(df, 'code')
                open = float(getattr(df, 'open'))
                close = float(getattr(df, 'price'))
                trade_date = getattr(df, 'date')

                dtend = datetime.datetime.strptime(trade_date, '%Y-%m-%d').date()
                dtbegin = dtend - datetime.timedelta(days=30)
                dtbeginstr = dtbegin.strftime('%Y-%m-%d')

                # 根据收盘价来折算筛选比例，收盘价越小比例越小
                k = 1
                if close < 10:
                    k = 0.2
                elif close < 20:
                    k = 0.4
                elif close < 30:
                    k = 0.5

                if open == 0:
                    continue

                #计算10日均线，并过滤收盘价低于10日均线的数据
                ts_code = stocklist[stocklist.symbol == symbol]['ts_code'].values[0]
                temp = pro.daily(ts_code=ts_code, start_date=dtbeginstr, end_date=trade_date)
                mean10 = temp.close.rolling(10).mean()
                mean10 = mean10.dropna()
                if len(mean10.values) > 0:
                    fmean10 = mean10.values[0]
                    if close < fmean10:
                        continue

                serials = pd.Series(
                    {'symbol': symbol, 'name': name, 'open': open, 'close': close, 'trade_date': trade_date}, name='')

                if (close - open)/open < det*k and close - open > 0:
                    positive = positive.append(serials)

                if (close - open)/open > -det*k and close - open < 0:
                    negative = negative.append(serials)

    positive.columns=['股票编码','股票名称','开盘价','收盘价','时间']
    negative.columns = ['股票编码', '股票名称', '开盘价','收盘价','时间']

    # (close-open)/open<det，导出文档名
    pfile = trade_date + 'positive.csv'
    # (close-open)/open>-det，导出文档名
    nfile = trade_date + 'negative.csv'

    positive.to_csv(pfile,encoding='gbk')
    negative.to_csv(nfile,encoding='gbk')

if __name__ == '__main__':
    # tushare注册后得到的token
    ts.set_token(config.token)
    # 交易日期
    trade_date = config.trade_date
    # 十字星范围，调整这个参数即可
    det = config.det
    # 获取数据
    getdata(det)